import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.*;
import java.util.Map.Entry;

public class ConsistentHash<T> {

    public static final int DEFAULT_NUMBER_OF_REPLICAS = 1000;

    private final SortedMap<Long, T> ring = new TreeMap<>();
    private final int numberOfReplicas;
    private MessageDigest md5;

    public ConsistentHash() {
        this(DEFAULT_NUMBER_OF_REPLICAS, new HashSet<>());
    }

    public ConsistentHash(int numberOfReplicas) {
        this(numberOfReplicas, new HashSet<>());
    }

    public ConsistentHash(int numberOfReplicas, Collection<T> nodes) {
        try {
            md5 = MessageDigest.getInstance("MD5");
        } catch (NoSuchAlgorithmException ignored) {
        }
        this.numberOfReplicas = numberOfReplicas;
        for (T node : nodes) {
            addNode(node);
        }
    }

    public void addNode(T node) {
        for (int i = 0; i < numberOfReplicas; i++) {
            ring.put(hash(virtualNodeKey(node, i)), node);
        }
    }

    public void addAll(Collection<T> nodes) {
        for (T node : nodes) {
            addNode(node);
        }
    }

    public void removeNode(T node) {
        for (int i = 0; i < numberOfReplicas; i++) {
            ring.remove(hash(virtualNodeKey(node, i)));
        }
    }

    public void removeAll() {
        ring.clear();
    }

    public boolean containsNode(T node) {
        for (int i = 0; i < numberOfReplicas; i++) {
            if (ring.containsKey(hash(virtualNodeKey(node, i)))) {
                return true;
            }
        }
        return false;
    }

    public T getNode(String objectKey) {
        if (ring.isEmpty()) {
            return null;
        }
        long hash = hash(objectKey);
        if (!ring.containsKey(hash)) {
            SortedMap<Long, T> tailMap = ring.tailMap(hash);
            hash = tailMap.isEmpty() ? ring.firstKey() : tailMap.firstKey();
        }
        return ring.get(hash);
    }

    public Collection<T> getAll() {
        return ring.values();
    }

    private String virtualNodeKey(T node, int replicaIndex) {
        return node.toString() + "-" + replicaIndex;
    }

    private long hash(String data) {
        md5.reset();
        md5.update(data.getBytes());
        byte[] digest = md5.digest();
        long h = 0;
        for (int i = 0; i < 4; i++) {
            h <<= 8;
            h |= ((int) digest[i]) & 0xFF;
        }
        return h;
    }

    public static void main(String[] args) {
        HashSet<String> nodes = new HashSet<String>();
        nodes.add("A");
        nodes.add("B");
        nodes.add("C");
        nodes.add("D");

        Map<String, Integer> stats = new HashMap<>();

        ConsistentHash<String> consistentHash = new ConsistentHash<>(1000, nodes);

        int count = 10000;
        for (int i = 0; i < count; i++) {
            String key = consistentHash.getNode(i + "");
            if (stats.containsKey(key)) {
                stats.put(consistentHash.getNode(i + ""), stats.get(key) + 1);
            } else {
                stats.put(consistentHash.getNode(i + ""), 1);
            }
        }
        showServer(stats);

        stats.clear();
        consistentHash.removeNode("A");
        System.out.println("------- remove A");
        for (int i = 0; i < count; i++) {
            String key = consistentHash.getNode(i + "");
            if (stats.containsKey(key)) {
                stats.put(consistentHash.getNode(i + ""), stats.get(key) + 1);
            } else {
                stats.put(consistentHash.getNode(i + ""), 1);
            }
        }
        showServer(stats);

        stats.clear();
        consistentHash.addNode("E");
        System.out.println("------- add E");
        for (int i = 0; i < count; i++) {
            String key = consistentHash.getNode(i + "");
            if (stats.containsKey(key)) {
                stats.put(consistentHash.getNode(i + ""), stats.get(key) + 1);
            } else {
                stats.put(consistentHash.getNode(i + ""), 1);
            }
        }
        showServer(stats);

        stats.clear();
        consistentHash.addNode("F");
        System.out.println("------- add F服务器  业务量加倍");
        count = count * 2;
        for (int i = 0; i < count; i++) {
            String key = consistentHash.getNode(i + "");
            if (stats.containsKey(key)) {
                stats.put(consistentHash.getNode(i + ""), stats.get(key) + 1);
            } else {
                stats.put(consistentHash.getNode(i + ""), 1);
            }
        }
        showServer(stats);
    }

    public static void showServer(Map<String, Integer> map) {
        for (Entry<String, Integer> m : map.entrySet()) {
            System.out.println("服务器 " + m.getKey() + "----" + m.getValue() + "个");
        }
    }
}